package me.shan.sort.heap;

import lombok.ToString;

import java.util.*;

/**
 * 加强堆，增加一个反向索引表，用以支持指定位置值更新O(logN)以及指定位置值删除O(logN)
 * <p>默认按照小根堆处理，comparator中小的排在父节点，大的排在子节点</p>
 * <p>元素类型E为基础数据类型时，由于反向索引表存在hash碰撞，因此需要将基础类再封装一层</p>
 *
 * @author liuchuang
 * @version 1.0
 * @since 2023/6/13 10:51
 */
@ToString
public class HeapEnhance<E> {

    /**
     * 列表实现的堆
     */
    private final List<E> heap;
    /**
     * 元素对应位置索引表
     */
    private final Map<E, Integer> elementIndexMap;
    /**
     * 比较器
     */
    private final Comparator<E> comparator;

    public HeapEnhance(Comparator<E> comparator) {
        this.heap = new ArrayList<>();
        this.elementIndexMap = new HashMap<>(32);
        this.comparator = comparator;
    }

    public void add(E value) {
        heap.add(value);
        elementIndexMap.put(value, heap.size() - 1);
        heapInsert(heap.size() - 1);
    }

    public E peek() {
        if (heap.size() == 0) {
            return null;
        } else {
            return heap.get(0);
        }
    }

    public E poll() {
        if (heap.size() == 0) {
            return null;
        } else {
            E value = heap.get(0);
            // 先交换，保证小根堆不会整个乱掉
            swap(0, heap.size() - 1);
            heap.remove(value);
            elementIndexMap.remove(value);
            heapify(0);
            return value;
        }
    }

    /**
     * 更新指定位置元素值
     *
     * @param index 待更新位置
     * @param value 更新之后的值
     */
    public void set(Integer index, E value) {
        if (index >= heap.size()) {
            return;
        }
        E oldValue = heap.get(index);
        if (value == oldValue) {
            return;
        }
        heap.set(index, value);
        elementIndexMap.remove(oldValue);
        elementIndexMap.put(value, index);
        heapInsert(index);
        heapify(index);
    }

    /**
     * 删除指定值
     *
     * @param value 待删除值
     */
    public void remove(E value) {
        if (elementIndexMap.containsKey(value)) {
            int index = elementIndexMap.get(value);
            int last = heap.size() - 1;
            if (index == last) {
                heap.remove(index);
                elementIndexMap.remove(value);
            } else {
                swap(index, last);
                heap.remove(last);
                elementIndexMap.remove(value);
                heapInsert(index);
                heapify(index);
            }
        }
    }

    private void heapInsert(int index) {
        int fatherIndex = (index - 1) / 2;
        while (comparator.compare(heap.get(index), heap.get(fatherIndex)) < 0) {
            swap(index, fatherIndex);
            index = fatherIndex;
            fatherIndex = (index - 1) / 2;
        }
    }

    private void heapify(int index) {
        int left = index * 2 + 1;
        while (left < heap.size()) {
            int least = left + 1 < heap.size() && comparator.compare(heap.get(left + 1), heap.get(left)) < 0 ? left + 1 : left;
            least = comparator.compare(heap.get(least), heap.get(index)) < 0 ? least : index;
            if (least == index) {
                break;
            }
            swap(index, least);
            index = least;
            left = index * 2 + 1;
        }
    }

    /**
     * 交付两个节点在堆中的值，以及反向索引表中的索引
     *
     * @param a 节点1
     * @param b 节点2
     */
    private void swap(int a, int b) {
        E valueA = heap.get(a);
        E valueB = heap.get(b);
        Integer indexA = elementIndexMap.get(valueA);
        Integer indexB = elementIndexMap.get(valueB);
        heap.set(indexA, valueB);
        heap.set(indexB, valueA);
        elementIndexMap.put(valueA, indexB);
        elementIndexMap.put(valueB, indexA);
    }

    public static void main(String[] args) {
        HeapEnhance<Wrapper<Integer>> heapEnhance = new HeapEnhance<>(Comparator.comparingInt(Wrapper::getValue));
        int len = 10;
        for (int i = len; i >= 0; i--) {
            heapEnhance.add(new Wrapper<>(i));
        }
        System.out.println(heapEnhance);
        heapEnhance.set(5, new Wrapper<>(4));
        System.out.println(heapEnhance);
        heapEnhance.remove(heapEnhance.heap.get(1));
        System.out.println(heapEnhance);
        System.out.println(heapEnhance.peek());
        System.out.println(heapEnhance.poll());
        System.out.println(heapEnhance);
    }

}

/**
 * 基础类型使用该类封装一层，确保hash以及equal能取到地址（防止hashMap中基础类型key重复）
 *
 * @param <P> 基础类型
 */
class Wrapper<P> {

    private final P value;

    public Wrapper(P value) {
        this.value = value;
    }

    public P getValue() {
        return value;
    }

    @Override
    public String toString() {
        return String.valueOf(this.value);
    }
}
